import time, redis, redis, os, json


def query_redis_link():
    BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    config_file = os.path.join(BASE_DIR, 'SaaSDeveMySqlConfig', 'database.json')
    with open(config_file) as f:
        database_config = json.load(f)
    return database_config["redis_link"]


class Snowflake:
    """雪花算法生成唯一id"""

    def __init__(self, redis_link):
        self.worker_id = 0  # 默认的worker_id为0
        self.timestamp_bits = 41  # 时间戳位数
        self.worker_id_bits = 10  # worker_id位数
        self.sequence_bits = 12  # 序列号位数
        self.max_worker_id = -1 ^ (-1 << self.worker_id_bits)  # worker_id的最大值
        self.max_sequence = -1 ^ (-1 << self.sequence_bits)  # 序列号的最大值
        self.sequence = 0  # 当前序列号
        self.last_timestamp = -1  # 上次生成ID的时间戳
        self.redis_client = redis.Redis(
            host=redis_link["HOST"],
            port=redis_link["PORT"],
            db=3,
            password=redis_link["PASSWORD"]
        )  # 连接Redis数据库
    def generate_id(self):
        timestamp = self.get_timestamp()  # 获取当前时间戳

        if timestamp < self.last_timestamp:  # 如果当前时间戳小于上次生成ID的时间戳，抛出异常
            raise Exception("当前时间戳小于上次生成ID")

        if timestamp == self.last_timestamp:  # 如果当前时间戳与上次生成ID的时间戳相同
            self.sequence = (self.sequence + 1) & self.max_sequence  # 序列号递增
            if self.sequence == 0:  # 如果序列号达到最大值，等待下一毫秒
                timestamp = self.wait_next_millis(self.last_timestamp)
        else:
            self.sequence = 0  # 时间戳发生变化，重置序列号为0

        self.last_timestamp = timestamp  # 更新上次生成ID的时间戳

        id_value = ((timestamp << (self.worker_id_bits + self.sequence_bits)) |  # 生成ID
                    (self.worker_id << self.sequence_bits) |
                    self.sequence)

        # 将ID和时间戳存储到Redis有序集合中
        self.redis_client.zadd('snowflake_ids', {id_value: timestamp})

        return id_value

    def get_timestamp(self):
        return int(time.time() * 1000)  # 获取当前的时间戳（精确到毫秒）

    def wait_next_millis(self, last_timestamp):
        timestamp = self.get_timestamp()
        while timestamp <= last_timestamp:  # 循环等待下一毫秒
            timestamp = self.get_timestamp()
        return timestamp

    def check_clock_drift(self, max_clock_drift):
        # 获取有序集合中最早的成员（即最小分值）
        earliest_member = self.redis_client.zrange('snowflake_ids', 0, 0, withscores=True)
        if earliest_member:
            earliest_timestamp = earliest_member[0][1]
            current_timestamp = self.get_timestamp()
            clock_drift = current_timestamp - earliest_timestamp
            if clock_drift > max_clock_drift:
                raise Exception("Clock drift detected. Difference: {} ms".format(clock_drift))


# 示例用法
snowflake = Snowflake(query_redis_link())
